"""
20_chunking_strategies.py
=========================
Document Chunking Strategies for Graph RAG

This module covers:
1. Fixed-size chunking
2. Semantic chunking
3. Recursive chunking
4. Document-based chunking
5. Sentence-based chunking
6. Overlap strategies
7. Storing chunks in Neo4j

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional, Callable
from dataclasses import dataclass, field
from enum import Enum
import re
import hashlib
from datetime import datetime
from neo4j import Driver
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class ChunkingStrategy(Enum):
    """Available chunking strategies."""
    FIXED_SIZE = "fixed_size"
    SENTENCE = "sentence"
    PARAGRAPH = "paragraph"
    RECURSIVE = "recursive"
    SEMANTIC = "semantic"


@dataclass
class TextChunk:
    """Represents a text chunk."""
    id: str
    content: str
    index: int
    start_char: int
    end_char: int
    metadata: Dict[str, Any] = field(default_factory=dict)
    embedding: List[float] = None
    
    @property
    def length(self) -> int:
        return len(self.content)
    
    @property
    def word_count(self) -> int:
        return len(self.content.split())


@dataclass
class Document:
    """Represents a document to be chunked."""
    id: str
    content: str
    title: str = ""
    source: str = ""
    doc_type: str = "text"
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)


class TextChunker:
    """
    Implements various text chunking strategies.
    
    Chunking is critical for RAG systems:
    - Too large: Dilutes relevant information
    - Too small: Loses context
    - No overlap: Misses boundary information
    """
    
    def __init__(
        self,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        length_function: Callable[[str], int] = len
    ):
        """
        Initialize chunker.
        
        Args:
            chunk_size: Target chunk size (characters or tokens)
            chunk_overlap: Overlap between chunks
            length_function: Function to measure text length
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.length_function = length_function
    
    def _generate_chunk_id(self, content: str, index: int) -> str:
        """Generate unique chunk ID."""
        hash_input = f"{content[:100]}{index}".encode()
        return hashlib.md5(hash_input).hexdigest()[:12]
    
    # =========================================================
    # FIXED-SIZE CHUNKING
    # =========================================================
    
    def fixed_size_chunk(
        self,
        text: str,
        doc_id: str = None
    ) -> List[TextChunk]:
        """
        Split text into fixed-size chunks with overlap.
        
        This is the simplest strategy. Works well for uniform text.
        
        Args:
            text: Text to chunk
            doc_id: Document ID for metadata
            
        Returns:
            List of TextChunk objects
            
        Example:
            text = "The quick brown fox jumps over the lazy dog."
            chunk_size = 20, overlap = 5
            
            Chunks:
            1. "The quick brown fox " (0-20)
            2. " fox jumps over the" (15-35)  # 5 char overlap
            3. " the lazy dog."      (30-44)
        """
        chunks = []
        start = 0
        index = 0
        
        while start < len(text):
            # Calculate end position
            end = start + self.chunk_size
            
            # Don't cut in middle of word if possible
            if end < len(text):
                # Look for word boundary
                space_pos = text.rfind(' ', start, end)
                if space_pos > start:
                    end = space_pos + 1
            
            chunk_text = text[start:end].strip()
            
            if chunk_text:
                chunk = TextChunk(
                    id=self._generate_chunk_id(chunk_text, index),
                    content=chunk_text,
                    index=index,
                    start_char=start,
                    end_char=end,
                    metadata={
                        "doc_id": doc_id,
                        "strategy": "fixed_size",
                        "chunk_size": self.chunk_size,
                        "overlap": self.chunk_overlap
                    }
                )
                chunks.append(chunk)
                index += 1
            
            # Move start position with overlap
            start = end - self.chunk_overlap
            if start >= len(text):
                break
        
        return chunks
    
    # =========================================================
    # SENTENCE-BASED CHUNKING
    # =========================================================
    
    def sentence_chunk(
        self,
        text: str,
        doc_id: str = None,
        sentences_per_chunk: int = 3
    ) -> List[TextChunk]:
        """
        Split text into chunks by sentences.
        
        Respects sentence boundaries, preserving semantic units.
        
        Args:
            text: Text to chunk
            doc_id: Document ID
            sentences_per_chunk: Number of sentences per chunk
            
        Returns:
            List of TextChunk objects
        """
        # Split into sentences
        sentence_pattern = r'(?<=[.!?])\s+'
        sentences = re.split(sentence_pattern, text)
        sentences = [s.strip() for s in sentences if s.strip()]
        
        chunks = []
        index = 0
        start_char = 0
        
        for i in range(0, len(sentences), sentences_per_chunk):
            chunk_sentences = sentences[i:i + sentences_per_chunk]
            chunk_text = ' '.join(chunk_sentences)
            
            # Find actual position in original text
            end_char = start_char + len(chunk_text)
            
            chunk = TextChunk(
                id=self._generate_chunk_id(chunk_text, index),
                content=chunk_text,
                index=index,
                start_char=start_char,
                end_char=end_char,
                metadata={
                    "doc_id": doc_id,
                    "strategy": "sentence",
                    "sentence_count": len(chunk_sentences),
                    "first_sentence": chunk_sentences[0][:50] + "..."
                }
            )
            chunks.append(chunk)
            
            start_char = end_char + 1
            index += 1
        
        return chunks
    
    # =========================================================
    # PARAGRAPH-BASED CHUNKING
    # =========================================================
    
    def paragraph_chunk(
        self,
        text: str,
        doc_id: str = None,
        max_chunk_size: int = None
    ) -> List[TextChunk]:
        """
        Split text by paragraphs, combining small ones.
        
        Preserves document structure while managing chunk size.
        
        Args:
            text: Text to chunk
            doc_id: Document ID
            max_chunk_size: Maximum chunk size (combines small paragraphs)
            
        Returns:
            List of TextChunk objects
        """
        max_chunk_size = max_chunk_size or self.chunk_size
        
        # Split by double newlines (paragraphs)
        paragraphs = re.split(r'\n\s*\n', text)
        paragraphs = [p.strip() for p in paragraphs if p.strip()]
        
        chunks = []
        current_chunk = []
        current_length = 0
        index = 0
        start_char = 0
        
        for para in paragraphs:
            para_length = len(para)
            
            # If adding this paragraph exceeds max size, save current chunk
            if current_length + para_length > max_chunk_size and current_chunk:
                chunk_text = '\n\n'.join(current_chunk)
                chunk = TextChunk(
                    id=self._generate_chunk_id(chunk_text, index),
                    content=chunk_text,
                    index=index,
                    start_char=start_char,
                    end_char=start_char + len(chunk_text),
                    metadata={
                        "doc_id": doc_id,
                        "strategy": "paragraph",
                        "paragraph_count": len(current_chunk)
                    }
                )
                chunks.append(chunk)
                
                start_char += len(chunk_text) + 2
                index += 1
                current_chunk = []
                current_length = 0
            
            current_chunk.append(para)
            current_length += para_length
        
        # Don't forget the last chunk
        if current_chunk:
            chunk_text = '\n\n'.join(current_chunk)
            chunk = TextChunk(
                id=self._generate_chunk_id(chunk_text, index),
                content=chunk_text,
                index=index,
                start_char=start_char,
                end_char=start_char + len(chunk_text),
                metadata={
                    "doc_id": doc_id,
                    "strategy": "paragraph",
                    "paragraph_count": len(current_chunk)
                }
            )
            chunks.append(chunk)
        
        return chunks
    
    # =========================================================
    # RECURSIVE CHUNKING
    # =========================================================
    
    def recursive_chunk(
        self,
        text: str,
        doc_id: str = None,
        separators: List[str] = None
    ) -> List[TextChunk]:
        """
        Recursively split text using hierarchy of separators.
        
        Tries to split at natural boundaries first, then falls back
        to smaller units if chunks are still too large.
        
        Args:
            text: Text to chunk
            doc_id: Document ID
            separators: List of separators in order of preference
            
        Returns:
            List of TextChunk objects
            
        Default separator hierarchy:
        1. Double newlines (paragraphs)
        2. Single newlines
        3. Periods (sentences)
        4. Spaces (words)
        5. Characters
        """
        separators = separators or ["\n\n", "\n", ". ", " ", ""]
        
        def _split_recursive(text: str, separators: List[str]) -> List[str]:
            """Recursively split text."""
            if not separators:
                return [text]
            
            separator = separators[0]
            remaining_separators = separators[1:]
            
            if separator == "":
                # Character-level split
                return [text[i:i + self.chunk_size] 
                        for i in range(0, len(text), self.chunk_size)]
            
            splits = text.split(separator)
            
            result = []
            current = []
            current_length = 0
            
            for split in splits:
                split_length = len(split)
                
                if current_length + split_length > self.chunk_size:
                    if current:
                        merged = separator.join(current)
                        if len(merged) > self.chunk_size:
                            # Still too big, recurse
                            result.extend(_split_recursive(merged, remaining_separators))
                        else:
                            result.append(merged)
                        current = []
                        current_length = 0
                
                current.append(split)
                current_length += split_length + len(separator)
            
            if current:
                merged = separator.join(current)
                if len(merged) > self.chunk_size and remaining_separators:
                    result.extend(_split_recursive(merged, remaining_separators))
                else:
                    result.append(merged)
            
            return result
        
        # Get recursive splits
        splits = _split_recursive(text, separators)
        
        # Convert to TextChunk objects
        chunks = []
        start_char = 0
        
        for index, chunk_text in enumerate(splits):
            chunk_text = chunk_text.strip()
            if chunk_text:
                chunk = TextChunk(
                    id=self._generate_chunk_id(chunk_text, index),
                    content=chunk_text,
                    index=index,
                    start_char=start_char,
                    end_char=start_char + len(chunk_text),
                    metadata={
                        "doc_id": doc_id,
                        "strategy": "recursive"
                    }
                )
                chunks.append(chunk)
                start_char += len(chunk_text) + 1
        
        return chunks
    
    # =========================================================
    # SEMANTIC CHUNKING (EMBEDDING-BASED)
    # =========================================================
    
    def semantic_chunk(
        self,
        text: str,
        doc_id: str = None,
        embedding_function: Callable[[str], List[float]] = None,
        similarity_threshold: float = 0.5
    ) -> List[TextChunk]:
        """
        Split text based on semantic similarity between sentences.
        
        Groups sentences that are semantically similar together.
        Requires an embedding function.
        
        Args:
            text: Text to chunk
            doc_id: Document ID
            embedding_function: Function to generate embeddings
            similarity_threshold: Threshold for grouping (0-1)
            
        Returns:
            List of TextChunk objects
        """
        if embedding_function is None:
            # Fallback to sentence chunking
            return self.sentence_chunk(text, doc_id)
        
        # Split into sentences
        sentence_pattern = r'(?<=[.!?])\s+'
        sentences = re.split(sentence_pattern, text)
        sentences = [s.strip() for s in sentences if s.strip()]
        
        if len(sentences) <= 1:
            return [TextChunk(
                id=self._generate_chunk_id(text, 0),
                content=text,
                index=0,
                start_char=0,
                end_char=len(text),
                metadata={"doc_id": doc_id, "strategy": "semantic"}
            )]
        
        # Get embeddings for each sentence
        embeddings = [embedding_function(s) for s in sentences]
        
        # Calculate cosine similarity between adjacent sentences
        def cosine_similarity(a: List[float], b: List[float]) -> float:
            import math
            dot = sum(x * y for x, y in zip(a, b))
            norm_a = math.sqrt(sum(x * x for x in a))
            norm_b = math.sqrt(sum(x * x for x in b))
            return dot / (norm_a * norm_b) if norm_a and norm_b else 0
        
        # Find split points (low similarity between adjacent sentences)
        split_points = [0]
        for i in range(1, len(sentences)):
            similarity = cosine_similarity(embeddings[i-1], embeddings[i])
            if similarity < similarity_threshold:
                split_points.append(i)
        split_points.append(len(sentences))
        
        # Create chunks
        chunks = []
        start_char = 0
        
        for idx in range(len(split_points) - 1):
            start_sent = split_points[idx]
            end_sent = split_points[idx + 1]
            
            chunk_sentences = sentences[start_sent:end_sent]
            chunk_text = ' '.join(chunk_sentences)
            
            chunk = TextChunk(
                id=self._generate_chunk_id(chunk_text, idx),
                content=chunk_text,
                index=idx,
                start_char=start_char,
                end_char=start_char + len(chunk_text),
                metadata={
                    "doc_id": doc_id,
                    "strategy": "semantic",
                    "sentence_count": len(chunk_sentences)
                }
            )
            chunks.append(chunk)
            start_char += len(chunk_text) + 1
        
        return chunks


class ChunkStore:
    """
    Stores and manages chunks in Neo4j.
    """
    
    def __init__(self, driver: Driver):
        self.driver = driver
    
    def store_document(self, document: Document) -> str:
        """
        Store a document node.
        
        Args:
            document: Document to store
            
        Returns:
            Document ID
        """
        query = """
        MERGE (d:Document {id: $id})
        SET d.title = $title,
            d.source = $source,
            d.docType = $doc_type,
            d.createdAt = datetime($created_at),
            d.contentLength = $content_length
        SET d += $metadata
        RETURN d.id as id
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                id=document.id,
                title=document.title,
                source=document.source,
                doc_type=document.doc_type,
                created_at=document.created_at.isoformat(),
                content_length=len(document.content),
                metadata=document.metadata
            )
            return result.single()["id"]
    
    def store_chunks(
        self,
        chunks: List[TextChunk],
        document_id: str
    ) -> int:
        """
        Store chunks and link to document.
        
        Creates:
        - Chunk nodes with content and metadata
        - HAS_CHUNK relationships from Document
        - NEXT relationships between sequential chunks
        
        Args:
            chunks: List of TextChunk objects
            document_id: Parent document ID
            
        Returns:
            Number of chunks stored
        """
        # Create chunks with relationships
        query = """
        MATCH (d:Document {id: $doc_id})
        UNWIND $chunks AS chunk
        MERGE (c:Chunk {id: chunk.id})
        SET c.content = chunk.content,
            c.index = chunk.index,
            c.startChar = chunk.start_char,
            c.endChar = chunk.end_char,
            c.wordCount = chunk.word_count,
            c.charCount = chunk.char_count
        SET c += chunk.metadata
        MERGE (d)-[:HAS_CHUNK]->(c)
        WITH c, chunk
        ORDER BY chunk.index
        WITH collect(c) AS chunks
        UNWIND range(0, size(chunks)-2) AS i
        WITH chunks[i] AS c1, chunks[i+1] AS c2
        MERGE (c1)-[:NEXT]->(c2)
        """
        
        chunk_data = [
            {
                "id": c.id,
                "content": c.content,
                "index": c.index,
                "start_char": c.start_char,
                "end_char": c.end_char,
                "word_count": c.word_count,
                "char_count": c.length,
                "metadata": c.metadata
            }
            for c in chunks
        ]
        
        with self.driver.session() as session:
            session.run(query, doc_id=document_id, chunks=chunk_data)
        
        return len(chunks)
    
    def store_chunk_embeddings(
        self,
        chunk_ids: List[str],
        embeddings: List[List[float]]
    ) -> int:
        """
        Store embeddings for chunks.
        
        Args:
            chunk_ids: List of chunk IDs
            embeddings: Corresponding embeddings
            
        Returns:
            Number of embeddings stored
        """
        query = """
        UNWIND $data AS item
        MATCH (c:Chunk {id: item.id})
        SET c.embedding = item.embedding
        """
        
        data = [
            {"id": cid, "embedding": emb}
            for cid, emb in zip(chunk_ids, embeddings)
        ]
        
        with self.driver.session() as session:
            session.run(query, data=data)
        
        return len(data)
    
    def get_chunks_by_document(
        self,
        document_id: str
    ) -> List[Dict[str, Any]]:
        """
        Get all chunks for a document.
        
        Args:
            document_id: Document ID
            
        Returns:
            List of chunk data
        """
        query = """
        MATCH (d:Document {id: $doc_id})-[:HAS_CHUNK]->(c:Chunk)
        RETURN c
        ORDER BY c.index
        """
        
        with self.driver.session() as session:
            result = session.run(query, doc_id=document_id)
            return [dict(record["c"]) for record in result]


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Document Chunking Strategies - Demonstration")
    print("=" * 60)
    
    # Sample medical document
    sample_text = """
    Diabetes mellitus is a chronic metabolic disorder characterized by elevated blood glucose levels. 
    The condition affects millions of people worldwide and can lead to serious health complications 
    if not properly managed.
    
    Type 2 diabetes is the most common form, accounting for approximately 90-95% of all cases. 
    It typically develops in adults over 45, but is increasingly seen in younger populations. 
    Risk factors include obesity, sedentary lifestyle, and family history.
    
    Treatment approaches vary based on individual patient needs. Initial management often involves 
    lifestyle modifications including diet changes and increased physical activity. Metformin is 
    typically the first-line medication for most patients.
    
    Regular monitoring of blood glucose levels is essential. Patients should also undergo periodic 
    screening for complications including retinopathy, nephropathy, and neuropathy. 
    Cardiovascular risk assessment is crucial given the increased risk of heart disease.
    
    Recent advances in treatment include GLP-1 receptor agonists and SGLT2 inhibitors, which offer 
    additional benefits beyond glucose control. These medications have shown cardiovascular and 
    renal protective effects in clinical trials.
    """
    
    chunker = TextChunker(chunk_size=300, chunk_overlap=50)
    
    # Fixed-size chunking
    print("\n📌 Fixed-Size Chunking (300 chars, 50 overlap)")
    print("-" * 40)
    fixed_chunks = chunker.fixed_size_chunk(sample_text, "DOC001")
    for chunk in fixed_chunks:
        print(f"\n   Chunk {chunk.index}: {chunk.length} chars")
        print(f"   Preview: {chunk.content[:80]}...")
    
    # Sentence chunking
    print("\n📌 Sentence-Based Chunking (3 sentences per chunk)")
    print("-" * 40)
    sentence_chunks = chunker.sentence_chunk(sample_text, "DOC001", sentences_per_chunk=3)
    for chunk in sentence_chunks:
        print(f"\n   Chunk {chunk.index}: {chunk.word_count} words")
        print(f"   Preview: {chunk.content[:80]}...")
    
    # Paragraph chunking
    print("\n📌 Paragraph-Based Chunking")
    print("-" * 40)
    para_chunks = chunker.paragraph_chunk(sample_text, "DOC001", max_chunk_size=500)
    for chunk in para_chunks:
        print(f"\n   Chunk {chunk.index}: {chunk.metadata.get('paragraph_count', 1)} paragraphs")
        print(f"   Preview: {chunk.content[:80]}...")
    
    # Recursive chunking
    print("\n📌 Recursive Chunking")
    print("-" * 40)
    recursive_chunks = chunker.recursive_chunk(sample_text, "DOC001")
    for chunk in recursive_chunks:
        print(f"\n   Chunk {chunk.index}: {chunk.length} chars")
        print(f"   Preview: {chunk.content[:80]}...")
    
    # Store in Neo4j (if available)
    print("\n📌 Storing Chunks in Neo4j")
    print("-" * 40)
    try:
        driver = get_driver()
        store = ChunkStore(driver)
        
        # Create document
        doc = Document(
            id="DOC001",
            content=sample_text,
            title="Diabetes Management Overview",
            source="medical_knowledge_base",
            doc_type="clinical_guide"
        )
        
        store.store_document(doc)
        stored = store.store_chunks(fixed_chunks, doc.id)
        print(f"   Stored {stored} chunks for document {doc.id}")
        
        # Retrieve
        retrieved = store.get_chunks_by_document(doc.id)
        print(f"   Retrieved {len(retrieved)} chunks")
        
        driver.close()
    except Exception as e:
        print(f"   Neo4j not available: {e}")
    
    print("\n" + "=" * 60)
    print("✅ Chunking demonstration complete!")
    print("=" * 60)
